#!/usr/bin/python
# Author: Jimmy Saw
# Usage: python blast_lineage_votes2.py blastfile.tbl 50
# 50 is max number of top hits to return
# It is supposed to give you the parent phylum of the hit organism from a BLAST result
# Needs to be in tabular format parsed by Blast2table Perl script
# Last argument can take either phylum or class for now
# Usage: python blast_lineage_votes2.py blastfile.blastp.tbl 20 phylum

import sys
import re
from Bio import Entrez
from Bio import SeqIO

def fmt(f):
    st = '{0:.4}'.format(f)
    return st

Entrez.email = "jimmysaw@gmail.com"

blastfile = sys.argv[1]
bf = open(blastfile, "rU")
bfl = bf.readlines()

si = blastfile.split('.')
seqid = str(si[0])

maxorgs = sys.argv[2]

choice = str(sys.argv[3])

hit_acc = re.compile('\w+\|(\S+)\|')
orgname = re.compile('(\w+) (\w+)')

#p = seqid.match(blastfile)
#bfn = p.group(1)

num = len(bfl)

maximum = 0

if num > int(maxorgs):
    maximum = int(maxorgs)
else:
    maximum = num

phyla = []

if len(bfl) == 0:
    #print bfn + "\t" + "no BLAST hit"
    pass
else:
    for line in bfl[0:maximum]:
        s = line.split('\t')
        evalue = s[1]
        identities = s[3]
        similarities = s[4]
        locus_tag = s[5]
        accession = s[9]
        desc = s[13].rstrip()
        searchterm = ""
        taxo = ""
        pattern = hit_acc.match(accession)
        acc = pattern.group(1)
        handle = Entrez.efetch(db="nucleotide", id=acc, rettype="gb")
        rec = SeqIO.read(handle, "genbank")
        if choice == 'phylum':
            taxo = rec.annotations["taxonomy"][1]
        elif choice == 'class':
            taxo = rec.annotations["taxonomy"][2]
        #print locus_tag + "\t" + taxo 
        phyla.append(taxo)

pls = len(phyla)

if pls > 0:
    unique_phyla = set(phyla)
    unique_phyla_list = list(unique_phyla)
    descriptions = []
    name_votes = []
    for d in unique_phyla_list:
        descriptions.append(d)
        count = phyla.count(d)
        name_votes.append(count)
    smallest = min(name_votes)
    largest = max(name_votes)
    max_index = name_votes.index(largest)
    min_index = name_votes.index(smallest)
    print seqid + "\t" + "Largest vote = " + str(largest) + "/" + str(maximum)\
        + " (" + str(fmt((float(largest)/float(maximum))*100)) + "%)" + "\t" + descriptions[max_index]
    #print "Smallest vote = " + str(smallest) + "\t" + descriptions[min_index]
else:
    print seqid + "\t" + "No BLAST hit" + "\t" + "-"

bf.close()
